import numpy as np
import torch
import datetime
import os
import json
from utils.reranking import re_ranking

json_save_path = '../json/resnet101_ibn_b/'
qfeat_path = './probFea_ibn_b.npy'
gfeat_path = './galFea_ibn_b.npy'
def comput():
    q_path = np.load('./q_path.npy')
    g_path = np.load('./g_path.npy')
    proFeat = np.load(qfeat_path)
    galFeat = np.load(gfeat_path)
    qf = torch.from_numpy(proFeat).cuda()
    gf = torch.from_numpy(galFeat).cuda()

    reranking_parameter = [10, 3, 0.6]
    max_rank = 200
    print('=> Enter reranking')
    print('k1={}, k2={}, lambda_value={}'.format(reranking_parameter[0], reranking_parameter[1],
                                                 reranking_parameter[2]))
    distmat = re_ranking(qf, gf, k1=reranking_parameter[0], k2=reranking_parameter[1],
                         lambda_value=reranking_parameter[2])
    # print(distmat,'distmat')
    num_q, num_g = distmat.shape
    indices = np.argsort(distmat, axis=1)
    data = dict()
    # print(len(q_path), 'self.img_name_q')
    # print(len(g_path),'self.img_name_g')
    for q_idx in range(num_q):
        order = indices[q_idx]  # select one row
        result_query = np.array(g_path)[order[:max_rank]]
        data[q_path[q_idx]] = [str(i) for i in result_query]
    return data, distmat, q_path, g_path

subfix = 2
DISTMAT_PATH = os.path.join(json_save_path, "distmat_{}.npy".format(subfix))
QUERY_PATH = os.path.join(json_save_path, "query_path_{}.npy".format(subfix))
GALLERY_PATH = os.path.join(json_save_path, "gallery_path_{}.npy".format(subfix))
if __name__ == '__main__':

    data, distmat, img_name_q, img_name_g = comput()
    np.save(DISTMAT_PATH, distmat)
    np.save(QUERY_PATH, img_name_q)
    np.save(GALLERY_PATH, img_name_g)

    data_all = {**data}
    nowTime = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
    with open(os.path.join(json_save_path, 'result_{}.json'.format(nowTime)), 'w', encoding='utf-8') as fp:
        json.dump(data_all, fp)